#include "LibNetImpl/UNetOnnx.h"
#include <opencv2/opencv.hpp>
#include <exception>

#include <QDir>
#include <QFile>
#include <QFileInfo>
#include <QCoreApplication>
#include <QCommandLineOption>
#include <QCommandLineParser>

#if defined(_MSC_VER)
#define CLOCKS_PER  CLOCKS_PER_SEC
#else
#define CLOCKS_PER  ((__clock_t) 10000000)
#endif

// https://www.cnblogs.com/findumars/p/6375644.html
// sample input: --model_path aa --input_mat bb --output_mat cc
// --model_path sample_input/model/v1_onnx.onnx --input_mat sample_input/input/2.png --output_mat sample_input/output/xxx/2.png

int main(int argc, char *argv[])
{
    // 1. get argument
    QCoreApplication app(argc, argv);
    QCommandLineParser parser;
    parser.addHelpOption();

    QCommandLineOption modelPath("model_path", "where_nnunet_model_path", "model_path");
    QCommandLineOption inputMat("input_mat", "input_image_from_path", "input_mat");
    QCommandLineOption outputMat("output_mat", "output_result_save_path", "output_mat");

    parser.addOption(modelPath);
    parser.addOption(inputMat);
    parser.addOption(outputMat);

    parser.process(app);

    if(!(parser.isSet(modelPath) && parser.isSet(inputMat) && parser.isSet(outputMat))){
        qCritical() << "Required argument must be set to run";
        exit(-2);
    }

    if(parser.value(modelPath).size() == 0 ||
       parser.value(inputMat).size() == 0 ||
       parser.value(outputMat).size() == 0){
        qCritical() << "Required argument can't be empty";
        exit(-2);
    }

    // 2. check data
    QString userModelPath = parser.value(modelPath);
    QString userInputPath = parser.value(inputMat);
    QString userOutputPath = parser.value(outputMat);

    QFileInfo modelInfo(userModelPath);
    if(!QFile::exists(userModelPath)){
        qCritical() << "Model Not Exists: " << userModelPath;
        exit(-3);
    }
    if(modelInfo.suffix() != "onnx"){
        qCritical() << "Model Not Available: " << userModelPath;
        exit(-3);
    }

    QFileInfo inputInfo(userInputPath);
    if(!QFile::exists(userInputPath)){
        qCritical() << "Input File Not Exists: " << userInputPath;
        exit(-3);
    }
    if(inputInfo.suffix() != "png"){
        qCritical() << "Input should be png file: " << userModelPath;
        exit(-3);
    }


    QFileInfo outInfo(userOutputPath);
    QString dirString = outInfo.absoluteDir().path();
    QDir outDir(dirString);
    if(!outDir.exists()){
        if(!outDir.mkpath(dirString)){
            qCritical() << "Output Dir Can't be create, check";
            exit(-3);
        }
    }

    try{
        std::string model = modelInfo.absoluteFilePath().toStdString();
        std::string input = inputInfo.absoluteFilePath().toStdString();
        std::string output = outInfo.absoluteFilePath().toStdString();
        double start, end;
        start = (double ) clock();
        UNetOnnx onnx(model);
        end = (double ) clock();
        qDebug() << "Time Cost: (init): " << (end - start) / CLOCKS_PER << " s";
        cv::Mat mat = onnx.infer(input);
        cv::imwrite(output, mat);
        end = (double ) clock();
        qDebug() << "Time Cost: (init+infer+save): " << (end - start) / CLOCKS_PER << " s";
    }catch (std::exception &e){
        qCritical() << "Error in infer model: " << e.what();
        exit(-4);
    }


}
